# Copyright 2025 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pandas as pd
import json
from pathlib import Path

def convert_parquet_to_json(parquet_file_path: str, json_file_path: str):
    """
    Converts a Parquet file generated by the previous script into a JSON file.
    Renames the 'context' column to 'input'.

    Args:
        parquet_file_path (str): The path to the input Parquet file.
        json_file_path (str): The desired path for the output JSON file.
    """
    # Load the Parquet file into a Pandas DataFrame
    df = pd.read_parquet(parquet_file_path)

    # Rename the 'context' column to 'input'
    records = []
    for i, record in enumerate(df.to_dict(orient='records')):
        records.append({
            'context': record['context'],
            'input': record['prompt'][0]['content'],
            'answers': list(record['reward_model']['ground_truth']),
            'num_docs': record['extra_info']['num_docs'],
            'index': record['extra_info']['index'],
        })


    # Save the list of dictionaries to a JSON file
    Path(json_file_path).parent.mkdir(parents=True, exist_ok=True)
    with open(json_file_path, 'w', encoding='utf-8') as f:
        json.dump(records, f, ensure_ascii=False, indent=4)

    print(f"Successfully converted '{parquet_file_path}' to '{json_file_path}'.")

if __name__ == "__main__":
    import os
    dataset_root = os.getenv("DATAROOT", "/mnt/hdfs/hongli/dataset/hotpotqa")
    input_parquet_file = f'{dataset_root}/hotpotqa_dev.parquet'
    output_json_file = f"{dataset_root}/eval_200.json"

    convert_parquet_to_json(input_parquet_file, output_json_file)
